package com.ti.demo;

import com.ti.demo.etl.IdTransformer;
import com.ti.demo.util.SslValidationUtil;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.model.transformer.SummaryMetadataEnricher;
import org.springframework.ai.reader.jsoup.JsoupDocumentReader;
import org.springframework.ai.reader.jsoup.config.JsoupDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;

import java.util.List;

@SpringBootApplication
public class SpringAiEtlApplication {

    public static final String DEFAULT_SUMMARY_EXTRACT_TEMPLATE = """
            Here is the content of the section:
            {context_str}

            Summarize the key topics and entities of the section in Chinese.

            Summary:""";

    public static void main(String[] args) {
        SpringApplication.run(SpringAiEtlApplication.class, args);
    }

    @Bean
    CommandLineRunner prepareData(DeepSeekChatModel aiClient, VectorStore vectorStore) {
        return new CommandLineRunner() {
            @Override
            public void run(String... args) throws Exception {

                String url = "https://data.eastmoney.com/report/info/AP202508191730497130.html";

                // 在创建 UrlResource 之前禁用 SSL 验证
                SslValidationUtil.disableSslVerification();

                // DocumentReader
                List<Document> htmlDocuments = new JsoupDocumentReader(url,
                        JsoupDocumentReaderConfig.builder()
                                .selector("body > div.main > div.main > div.zw-content > div.left > div.ctx-box")
                                .charset("UTF-8")  // Use ISO-8859-1 encoding
                                .includeLinkUrls(false) // Include link URLs in metadata
                                .additionalMetadata("source", url) // Add custom metadata
                                .build()).read();

//                List<Document> htmlDocuments = new TikaDocumentReader("https://data.eastmoney.com/report/info/AP202508191730583378.html").read();


                // DocumentTransformer
                // 有时候，我们希望 etl 任务能可重复执行，那么其 id 就需要和内容相关，不能随机生成
                htmlDocuments = new IdTransformer().apply(htmlDocuments);
                List<Document> summaryDocuments = new SummaryMetadataEnricher(aiClient,
                        List.of(SummaryMetadataEnricher.SummaryType.PREVIOUS, SummaryMetadataEnricher.SummaryType.CURRENT, SummaryMetadataEnricher.SummaryType.NEXT),
                        DEFAULT_SUMMARY_EXTRACT_TEMPLATE,
                        MetadataMode.ALL)
                        .apply(htmlDocuments);


                // DocumentWriter
                vectorStore.write(summaryDocuments);
            }
        };
    }

}
